import argparse
from pathlib import Path
import shutil

from ultralytics import YOLO
import torch


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Run YOLO inference using a trained model.")
    parser.add_argument(
        "--weights",
        type=str,
        default="weights/best.pt",
        help="Path to model weights (.pt or exported formats)",
    )
    parser.add_argument(
        "--source",
        type=str,
        default="sleep.jpg",
        help=(
            "Inference source: image/video path, directory, URL, or webcam index (e.g., 0)"
        ),
    )
    parser.add_argument(
        "--conf",
        type=float,
        default=0.6,
        help="Confidence threshold",
    )
    parser.add_argument(
        "--device",
        type=str,
        default="auto",
        help="Compute device: 'auto', 'cpu', 'mps', 0, 1, etc.",
    )
    parser.add_argument(
        "--show",
        action="store_true",
        help="Show inference windows",
    )
    parser.add_argument(
        "--save",
        action="store_true",
        help="Save annotated results to runs/ directory",
    )
    parser.add_argument(
        "--project",
        type=str,
        default="runs/detect",
        help="Project directory to save results",
    )
    parser.add_argument(
        "--name",
        type=str,
        default="predict",
        help="Run name (subdirectory under project)",
    )
    parser.add_argument(
        "--out",
        type=str,
        default="",
        help="Optional path to copy the single output image/video to",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()

    weights_path = Path(args.weights)
    if not weights_path.exists():
        raise FileNotFoundError(f"Weights not found: {weights_path}")

    model = YOLO(str(weights_path))

    # Resolve device 'auto' -> prefer Apple MPS, else CPU
    device = args.device
    if device == "auto":
        if torch.backends.mps.is_available():
            device = "mps"
        elif torch.cuda.is_available():
            device = "0"
        else:
            device = "cpu"

    # Ensure saving is enabled by default to produce an output file
    save_results = True if (args.save or True) else False

    # The Ultralytics API accepts many source types: path, dir, URL, int (webcam)
    model.predict(
        source=args.source,
        conf=args.conf,
        device=device,
        show=args.show,
        save=save_results,
        project=args.project,
        name=args.name,
        stream=False,
    )

    # Report saved files and optionally copy to --out
    save_dir = Path(args.project) / args.name
    if save_dir.exists():
        saved = sorted([p for p in save_dir.iterdir() if p.is_file()])
        if saved:
            print(f"Saved results to: {save_dir.resolve()}")
            for p in saved:
                print(f" - {p.name}")
            if args.out:
                out_path = Path(args.out)
                if len(saved) == 1:
                    out_path.parent.mkdir(parents=True, exist_ok=True)
                    shutil.copy2(saved[0], out_path)
                    print(f"Copied output to: {out_path.resolve()}")
                else:
                    print("--out specified but multiple files were saved; skipping copy.")


if __name__ == "__main__":
    main()


